import numpy as np
from PIL import Image
from torchvision import transforms
from pathlib import Path
from tqdm import tqdm
import os


def calculate_dataset_stats(folder_path, preprocess_transform, stop_after_n=None):
    """
    计算数据集中所有图片在Resize和Crop后的全局RGB均值和标准差。
    使用在线算法以节省内存。

    参数:
    folder_path (str): 包含图片A的文件夹路径。
    preprocess_transform (transforms.Compose): 用于统计的预处理 (Resize + Crop)。
    stop_after_n (int, optional): 仅计算前N张图片，用于快速测试。

    返回:
    (np.array, np.array): 全局均值 (3,) 和 全局标准差 (3,)
    """
    print(f"--- 1. 开始计算文件夹 '{folder_path}' 的全局颜色统计数据 ---")

    # 支持的图片格式
    image_extensions = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}

    # 使用 pathlib 递归查找所有图片
    image_paths = [
        p for p in Path(folder_path).rglob("*") if p.suffix.lower() in image_extensions
    ]

    if not image_paths:
        print(f"错误: 在 '{folder_path}' 中未找到任何图片。")
        return None, None

    if stop_after_n:
        print(f"!! 警告: 仅处理前 {stop_after_n} 张图片用于测试。")
        image_paths = image_paths[:stop_after_n]

    # 初始化累加器 (使用float64以保证精度)
    sum_pixels = np.zeros(3, dtype=np.float64)
    sum_sq_pixels = np.zeros(3, dtype=np.float64)
    total_pixel_count = 0

    num_images = len(image_paths)
    print(f"总共找到 {num_images} 张图片。")

    for img_path in tqdm(image_paths, desc="计算统计数据"):
        try:
            with Image.open(img_path).convert("RGB") as img:
                # 1. 应用 Resize(224) 和 CenterCrop(224)
                img_processed = preprocess_transform(img)

                # 2. 转换为Numpy
                arr = np.array(img_processed).astype(np.float64)

                # 3. 获取像素数量 (224 * 224)
                h, w, c = arr.shape
                pixel_count = h * w
                total_pixel_count += pixel_count

                # 4. 累加总和与平方和
                # arr.reshape(-1, 3) 将 (H, W, 3) 变为 (H*W, 3)
                sum_pixels += np.sum(arr.reshape(-1, 3), axis=0)
                sum_sq_pixels += np.sum(np.square(arr.reshape(-1, 3)), axis=0)

        except Exception as e:
            print(f"\n警告: 跳过图片 {img_path} (错误: {e})")

    if total_pixel_count == 0:
        print("错误: 没有成功处理任何图片。")
        return None, None

    # 5. 计算最终的均值和标准差
    # 均值: E[X] = sum(X) / N
    global_mean = sum_pixels / total_pixel_count

    # 标准差: std(X) = sqrt( E[X^2] - (E[X])^2 )
    # E[X^2] = sum(X^2) / N
    global_var = (sum_sq_pixels / total_pixel_count) - np.square(global_mean)
    # 确保方差非负（由于浮点数精度问题可能出现极小的负数）
    global_var[global_var < 0] = 0
    global_std = np.sqrt(global_var)

    # 6. 处理标准差为0的情况（例如，如果所有图片都是纯色）
    global_std[global_std == 0] = 1.0

    print("--- 统计数据计算完成 ---")
    print(f"全局均值 (RGB) (A): {global_mean}")
    print(f"全局标准差 (RGB) (A): {global_std}")

    return global_mean, global_std


def color_correct_with_stats(img_target, mean_ref, std_ref, preprocess_transform):
    """
    使用预先计算的全局统计数据 (mean_ref, std_ref) 来校正 img_target (B)。

    参数:
    img_target (PIL.Image): 需要被校正的图片B。
    mean_ref (np.array): 数据集A的全局均值。
    std_ref (np.array): 数据集A的全局标准差。
    preprocess_transform (transforms.Compose): 用于计算B的统计数据 (Resize + Crop)。

    返回:
    PIL.Image: 颜色校正后的图片B。
    """

    print("\n--- 2. 正在应用颜色校正到图片B ---")

    # 1. 计算 *图片B* 自身的统计数据 (在Crop后)
    img_target_stat = preprocess_transform(img_target)
    arr_target_stat = np.array(img_target_stat).astype(np.float64)

    mean_target = np.mean(arr_target_stat, axis=(0, 1))
    std_target = np.std(arr_target_stat, axis=(0, 1))
    std_target[std_target == 0] = 1  # 防止除以零

    print(f"图片B 均值 (RGB): {mean_target}")
    print(f"图片B 标准差 (RGB): {std_target}")

    # 2. 在 *原始* 目标图片上应用校正
    arr_target_original = np.array(img_target).astype(np.float64)

    # (B - mean_B) / std_B
    arr_corrected = arr_target_original - mean_target
    arr_corrected /= std_target

    # * std_A + mean_A
    arr_corrected = arr_corrected * std_ref + mean_ref

    # 3. 将值裁剪到有效的 [0, 255] 范围内
    arr_corrected = np.clip(arr_corrected, 0, 255)

    # 4. 转换回 uint8 并转为 PIL Image
    img_corrected = Image.fromarray(arr_corrected.astype(np.uint8))

    return img_corrected


# --- 主执行程序 ---
if __name__ == "__main__":
    # (取消注释) 实际使用时:
    # try:
    #     img_a = Image.open(
    #         "/data0/lcy/data/LNM/LNM_slices_uni_v1_processed/thumbnails/18-16971.jpg"
    #     ).convert("RGB")
    #     img_b = Image.open(
    #         "/data0/lcy/data/LNM/LNM_Zhujiang_uni_v1_processed/thumbnails/2f619d5b-e241-c214-ab82-3c2c1adca579_043338.jpg"
    #     ).convert("RGB")
    # except FileNotFoundError:
    #     print("未找到示例图片，将使用自动生成的色块。")
    # (保持上面的示例图片创建代码)
    # --- 1. 配置 ---

    # !! 修改这里: 指向您的训练集文件夹A
    PATH_TO_FOLDER_A = "/data0/lcy/data/LNM/LNM_slices_uni_v1_processed/thumbnails/"

    # !! 修改这里: 指向您想校正的图片B
    PATH_TO_IMAGE_B = "/data0/lcy/data/LNM/LNM_Zhujiang_uni_v1_processed/thumbnails/2f619d5b-e241-c214-ab82-3c2c1adca579_043338.jpg"

    # !! 修改这里: 校正后图片的保存路径
    PATH_TO_OUTPUT = "image_B_corrected.png"

    # (可选) 用于快速测试，设置 N
    # 例如: TEST_RUN_N_IMAGES = 10
    TEST_RUN_N_IMAGES = 30

    # --- 2. 创建示例文件 (仅用于测试) ---
    # (实际使用时您可以注释掉这个创建块)
    if not os.path.exists(PATH_TO_FOLDER_A):
        print(f"创建示例文件夹: {PATH_TO_FOLDER_A}")
        os.makedirs(PATH_TO_FOLDER_A, exist_ok=True)
        # 创建一些偏蓝色的 "训练集" 图片
        for i in range(5):
            noise = np.random.randint(-20, 20, (300, 300, 3))
            color = (70 + i * 5, 130 - i * 3, 180 + i * 2)
            img_a = Image.fromarray(
                np.clip(
                    np.array(Image.new("RGB", (300, 300), color)) + noise, 0, 255
                ).astype(np.uint8)
            )
            img_a.save(os.path.join(PATH_TO_FOLDER_A, f"train_img_{i}.png"))
        TEST_RUN_N_IMAGES = 5  # 强制在示例上使用

    if not os.path.exists(PATH_TO_IMAGE_B):
        print(f"创建示例图片: {PATH_TO_IMAGE_B}")
        # 创建一张偏红色的 "目标" 图片
        noise = np.random.randint(-30, 30, (400, 600, 3))
        img_b = Image.fromarray(
            np.clip(
                np.array(Image.new("RGB", (600, 400), color=(200, 100, 50))) + noise,
                0,
                255,
            ).astype(np.uint8)
        )
        img_b.save(PATH_TO_IMAGE_B)

    # --- 3. 执行 ---

    # 定义用于 *统计* 的预处理 (来自您的Resize和Crop)
    preprocess_for_stats = transforms.Compose(
        [
            transforms.Resize(224),
            transforms.CenterCrop(224),
        ]
    )

    try:
        # 1. 计算数据集A的全局统计数据
        mean_A, std_A = calculate_dataset_stats(
            PATH_TO_FOLDER_A, preprocess_for_stats, stop_after_n=TEST_RUN_N_IMAGES
        )

        if mean_A is not None:
            # 2. 加载图片B
            img_b = Image.open(PATH_TO_IMAGE_B).convert("RGB")

            # 3. 执行颜色校正
            img_b_corrected = color_correct_with_stats(
                img_b, mean_A, std_A, preprocess_for_stats
            )

            # 4. 保存结果
            img_b_corrected.save(PATH_TO_OUTPUT)

            print("\n--- 3. 完成 ---")
            print(f"原始图片B: {PATH_TO_IMAGE_B}")
            print(f"校正后图片已保存到: {PATH_TO_OUTPUT}")

    except FileNotFoundError as e:
        print(f"\n错误: 找不到文件或文件夹。请检查您的路径。")
        print(f"详细信息: {e}")
    except Exception as e:
        print(f"\n发生意外错误: {e}")
